package com.corpgovernment.common.utils;

import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.executor.BatchResult;
import org.apache.ibatis.session.ExecutorType;
import org.apache.ibatis.session.SqlSession;
import org.apache.ibatis.session.SqlSessionFactory;
import org.springframework.stereotype.Component;
import org.springframework.transaction.support.TransactionSynchronizationManager;


import javax.annotation.Resource;
import java.util.Arrays;
import java.util.List;
import java.util.function.BiFunction;


@Slf4j
@Component
public class MybatisBatchUtils {
 
    /**
    * 每次处理1000条
    */
    private static final int BATCH_SIZE = 1000;

    @Resource
    private SqlSessionFactory sqlSessionFactory;
 
    /**
    * 批量处理修改或者插入 [备注，批量更新的sql，如 8条更新 失败1条 不会主动触发Exception，需上游做出判定后主动发起回滚]
    * JDBC 驱动程序可以通过 executeBatch() 方法批量执行 SQL 语句并返回每条 SQL 语句执行后受影响的行数（成功条数）。该方法适用于批量执行 INSERT、UPDATE、DELETE 等 SQL 语句。
    * @param data     需要被处理的数据
    * @param mapperClass  Mybatis的Mapper类
    * @param function 自定义处理逻辑
    * @return int 影响的总行数
    */
    public <T,U,R> long batchUpdateOrInsert(List<T> data, Class<U> mapperClass, BiFunction<T,U,R> function) {
        int i = 1;
        long succ = 0;
        SqlSession batchSqlSession = sqlSessionFactory.openSession(ExecutorType.BATCH);
        try {
            U mapper = batchSqlSession.getMapper(mapperClass);
            int size = data.size();
            for (T element : data) {
                function.apply(element,mapper);
                if ((i % BATCH_SIZE == 0) || i == size) {
                    List<BatchResult> batchResultList = batchSqlSession.flushStatements();
                    for (int j = 0; j < batchResultList.size(); j++) {
                        int[] currBatch = batchResultList.get(j).getUpdateCounts();
                        long count = Arrays.stream(currBatch).filter(z -> z == 1).count();
                        succ = succ + count;
                    }
                }
                i++;
            }
            // 非事务环境下强制commit，事务情况下该commit相当于无效
            batchSqlSession.commit(!TransactionSynchronizationManager.isSynchronizationActive());
        } catch (Exception e) {
            batchSqlSession.rollback();
            succ = 0;
            log.error(e.getMessage(), e);
        } finally {
            batchSqlSession.close();
        }
        return succ;
    }
}